import datetime
from io import BytesIO
import io
import os
import uuid
import zipfile
from seal import EncryptionParameters, SEALContext, CoeffModulus, PlainModulus, scheme_type
from app.pkg.circuits import BooleanCircuit
from app.pkg.fhe_decryption import FileDecryptor
from app.pkg.fhe_encryption import FileEncryptor
from app.pkg.key_management import KeyManager
from app.pkg.paillier_handler import PaillierZKFileHandler
from app.pkg.paillier_zk import PaillierZK
from flask import Blueprint, current_app, request, jsonify, g, send_file
from app.common.tool.direction import get_filenames
from app.common.custom import CustomResponse
from werkzeug.utils import secure_filename

oss_routes = Blueprint('oss', __name__)
  
@oss_routes.route('/upload', methods=['POST'])
def update_oss():
    resp = CustomResponse('更新成功',200)
    try:
        encryType = request.headers.get('Encry-Type')
        # 获取上传的文件
        if 'file' not in request.files:
            return jsonify({"message": "没有检测到文件"}), 400

        file = request.files['file']
        # 确保文件是 ZIP 格式
        # if not file.filename.endswith('.zip'):
        #     return jsonify({"message": "只能上传 ZIP 文件"}), 400

        filename = secure_filename(file.filename)  # 确保文件名安全
        if encryType == '1':
            parms = EncryptionParameters(scheme_type.bgv)
            parms.set_poly_modulus_degree(8192)
            parms.set_coeff_modulus(CoeffModulus.BFVDefault(8192))  # 使用BFVDefault生成兼容链
            parms.set_plain_modulus(PlainModulus.Batching(8192, 20))
            context = SEALContext(parms)
            key_manager = KeyManager(g.OssHandler)
            encryptor = FileEncryptor(context, key_manager, g.OssHandler)
            encryptor.encrypt_file(file, filename, BooleanCircuit())
        elif encryType == '2': 
            zk = PaillierZK(512)
            handler = PaillierZKFileHandler(zk,g.OssHandler)
    
            handler.encrypt_file(file, filename)

        return jsonify(resp.to_dict()), 200

    except Exception as e:
        print(f"文件上传失败: {e}")
        return jsonify({"message": "服务器错误"}), 500

@oss_routes.route('/download', methods=['POST'])
def download_oss():
    resp = CustomResponse('更新成功',200)
    try:
        data = request.get_json()
        filename = data.get('filename')
        realname = data.get('realname')
        encryType = request.headers.get('Encry-Type')
        if encryType == '1':
            parms = EncryptionParameters(scheme_type.bgv)
            parms.set_poly_modulus_degree(8192)
            parms.set_coeff_modulus(CoeffModulus.BFVDefault(8192))  # 使用BFVDefault生成兼容链
            parms.set_plain_modulus(PlainModulus.Batching(8192, 20))
            context = SEALContext(parms)
            key_manager = KeyManager(g.OssHandler)
            decryptor = FileDecryptor(context, key_manager, g.OssHandler)
            data = decryptor.decrypt_file(f'{filename}/')
        elif encryType == '2': 
            zk = PaillierZK(512)
            handler = PaillierZKFileHandler(zk,g.OssHandler)
    
            data = handler.decrypt_file(filename)

        # 发送文件内容作为附件
        return send_file(
            io.BytesIO(data),
            mimetype='application/octet-stream',
            as_attachment=True,
            download_name=realname  # 可以改成 .jpg 或原始后缀
        )

    except Exception as e:
        print(f"文件下载失败: {e}")
        return jsonify({"message": "服务器错误"}), 500